{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7530c93e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install \"sagemaker>=2.48.0\" \"transformers==4.6.1\" \"datasets[s3]==1.6.2\" --upgrade -i https://opentuna.cn/pypi/web/simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "760ece84",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker.huggingface\n",
    "import sagemaker\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "role = sagemaker.get_execution_role()\n",
    "\n",
    "print(f\"IAM role arn used for running training: {role}\")\n",
    "print(f\"S3 bucket used for storing artifacts: {sess.default_bucket()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56036ad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "from datasets import load_dataset\n",
    "\n",
    "df_article_summary_full = pd.read_parquet('./meta_description.parquet', engine='pyarrow')\n",
    "df_article_summary_full[['original_text','meta_descrption']].to_csv('./total.csv',index=False)\n",
    "\n",
    "total_data = pd.read_csv('./total.csv')\n",
    "x = total_data[-total_data['meta_descrption'].isnull()]\n",
    "x.columns = ['article','summarization']\n",
    "\n",
    "# use csv file to test \n",
    "x[:1000].to_csv('./train.csv',index=False,encoding='utf-8')\n",
    "x[1000:1200].to_csv('./test.csv',index=False,encoding='utf-8')\n",
    "x[1200:1400].to_csv('./dev.csv',index=False,encoding='utf-8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "643b0598",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "prefix='hk01'\n",
    "\n",
    "bucket = sess.default_bucket() \n",
    "boto3.Session().resource(\"s3\").Bucket(bucket).Object(\n",
    "    os.path.join(prefix, \"train/train.csv\")\n",
    ").upload_file(\"./train.csv\")\n",
    "boto3.Session().resource(\"s3\").Bucket(bucket).Object(\n",
    "    os.path.join(prefix, \"test/test.csv\")\n",
    ").upload_file(\"./test.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6a911ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_input_path = f's3://{sess.default_bucket()}/{prefix}/train/train.csv'\n",
    "test_input_path = f's3://{sess.default_bucket()}/{prefix}/test/test.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09462081",
   "metadata": {},
   "outputs": [],
   "source": [
    "git_config = {'repo': 'https://gitee.com/whn09/transformers.git','branch': 'v4.6.1.1'} # v4.6.1 is referring to the `transformers_version` you use in the estimator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31a8a06",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparameters={'per_device_train_batch_size':1,\n",
    "                 'per_device_eval_batch_size': 1,\n",
    "                 'model_name_or_path': 'google/mt5-base',\n",
    "                 'train_file':'/opt/ml/input/data/train/train.csv',\n",
    "                 'validation_file':'/opt/ml/input/data/test/test.csv',\n",
    "                 'test_file':'/opt/ml/input/data/test/test.csv',\n",
    "                 'do_train': True,\n",
    "                 'do_predict': True,\n",
    "                 'do_eval': True,\n",
    "                 'text_column':'article',\n",
    "                 'summary_column':'summarization',\n",
    "                 'save_total_limit':3,\n",
    "                 'num_train_epochs': 1,\n",
    "                 'predict_with_generate': True,\n",
    "                 'output_dir': '/opt/ml/model',\n",
    "                 'num_train_epochs': 1,\n",
    "                 'learning_rate': 5e-5,\n",
    "                 'seed': 7,\n",
    "                 'fp16': False,\n",
    "                 'source_prefix': \"summarize: \",\n",
    "                 'eval_steps': 1000,\n",
    "                 }\n",
    "\n",
    "# configuration for running training on smdistributed Data Parallel\n",
    "#distribution = {'smdistributed':{'dataparallel':{ 'enabled': True }}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e46f657a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import HuggingFace\n",
    "\n",
    "# create the Estimator\n",
    "huggingface_estimator = HuggingFace(\n",
    "      entry_point='run_summarization.py', # script\n",
    "      source_dir='./examples/pytorch/summarization', # relative path to example\n",
    "      git_config=git_config,\n",
    "      instance_type='ml.p3.2xlarge', # here better to use ml.p3dn.24xlarge if available\n",
    "      instance_count=1,\n",
    "      volume_size=500,\n",
    "      transformers_version='4.6',\n",
    "      pytorch_version='1.7',\n",
    "      py_version='py36',\n",
    "      role=role,\n",
    "      base_job_name='mt5', \n",
    "      hyperparameters = hyperparameters\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0f76166",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "huggingface_estimator.fit({'train':training_input_path,'test':test_input_path})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "428a6537",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = huggingface_estimator.deploy(1,\"ml.g4dn.xlarge\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89d0beeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation = '''Jeff: Can I train a 🤗 Transformers model on Amazon SageMaker? \n",
    "    Philipp: Sure you can use the new Hugging Face Deep Learning Container. \n",
    "    Jeff: ok.\n",
    "    Jeff: and how can I get started? \n",
    "    Jeff: where can I find documentation? \n",
    "    Philipp: ok, ok you can find everything here. https://huggingface.co/blog/the-partnership-amazon-sagemaker-and-hugging-face                                           \n",
    "    '''\n",
    "\n",
    "data= {\"inputs\":conversation}\n",
    "\n",
    "predictor.predict(data)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_latest_p36",
   "language": "python",
   "name": "conda_pytorch_latest_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
